package bm25;

/**
 * create by lidongdong 18-9-8
 */

import java.io.*;
import java.nio.file.Paths;
import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.List;
import java.util.Map.Entry;

import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.similarities.BM25Similarity;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.RAMDirectory;


public class TrecModel {
    IndexReader reader = null;
    int resNum;
    String runName;
    String resPath;
    private static final int relevantNum = 20;
    private static final int irrelevantNum = 80;

    public TrecModel(int resNum, String runName) {
        this.resNum = resNum;
        this.runName = runName;
    }

    public static void main(String[] args) throws IOException, ParseException {
        // 所有的查询文件所在目录
        String docsPath = "/home/jack/train_data/trec";
        // 查询所在文件
        String topicsPath = "/home/jack/train_data/trec-query/query2015.txt";
        // query 文件和内容相关度文件
        String qrelsPath = "/home/jack/train_data/trec-query/qrels-sampleval-2015.txt";
        // 直接的建立索引
        String indexPath = "trec.index";
        String runName = "bm25_2015";
        int number = 1000;

        TrecModel tm = new TrecModel(number, runName);

        // 创建索引, 写的真的让人迷糊
        tm.makeIndex(docsPath, indexPath, true);
        float k1 = 1.5f, b = 0.8f;
        double[] k1s = TRECFloatRange(1, 4, 0.2);
        double[] bs = TRECFloatRange(0.75, 1, 0.05);
        Entry<Float, Float> params = tm.gridSearchBM25(topicsPath, qrelsPath, k1s, bs);
        k1 = params.getKey();
        b = params.getValue();

        String topics2015Path = "/home/jack/train_data/trec-query/query2015.txt";
        String qrels2015Path = "/home/jack/train_data/trec-query/qrels-sampleval-2015.txt";
        tm.setRunName("bm25_2015");

        // 进行评测
        tm.QueryTopicBM25(topics2015Path, k1, b, null);
        float ndcg = tm.infNDCGOfQueryResult(qrels2015Path, tm.getResPath());

        System.out.println("Topics 2015 (k1=" + k1 +",b=" + b +") : NDCG = " + ndcg );
    }

    /**
     * save all the file to index file
     * @param docsPath
     * @param indexPath
     * @param create
     */
    private void makeIndex(String docsPath, String indexPath, boolean create) throws IOException {
        File f = new File(indexPath);
        if (f.exists()){
            System.out.println(String.format("%s do exists", indexPath));
        }else {
            System.out.println(String.format("%s do not exists", indexPath));
            // 使用bm25simliarity来创建索引， 这里有什么区别么？？？
            if (docsPath != null) {
                IndexFile.IndexProcess(indexPath, docsPath, create);
            }
        }
        this.reader = DirectoryReader.open(new RAMDirectory(FSDirectory.open(Paths.get(indexPath)), IOContext.READONCE));
    }

    /**
     *
     * @param s
     * @param e
     * @param step
     * @return
     */
    private static double[] TRECFloatRange(double s, double e, double step) {
        double[] res = new double[((int)Math.floor((e - s)/step) + 1)];
        int idx = 0;
        for (; s <= e; s += step){
            res[idx++] = s;
        }
        return res;
    }

    public List<String> getTopics(String topicPath, boolean useSummary){
        File file = new File(topicPath);
        Long fileLength = file.length();
        byte[] content = new byte[fileLength.intValue()];
        try{
            FileInputStream inputStream = new FileInputStream(file);
            inputStream.read(content);
            inputStream.close();
            //将content转换为String list
            String str = new String(content);
            List<String> res = new ArrayList<>();
            String[] temp = str.split("\n");
            for (String t : temp){
                String discription = t.split("\t")[1];
                String summary = t.split("\t")[2];
                if(useSummary){
                    res.add(summary.substring(1, summary.length() - 1));
                }else {
                    res.add(discription.substring(1, discription.length() - 1));
                }
            }
            return res;
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 基于BM25模型, 对于每个topic检索出1000个文档
     * @param topicsPath
     * @param k1
     * @param b
     * @param gridResPath
     */
    private void QueryTopicBM25(String topicsPath, float k1, float b, String gridResPath) throws IOException, ParseException {
        // 使用 summary 作为检索的条件
        List<String> topics = getTopics(topicsPath, true);
        if (reader == null){
            System.out.println("The variable reader is not initialized");
            System.exit(0);
        }
        QueryParser queryParser = new QueryParser("contents", new StandardAnalyzer());
        IndexSearcher searcher = new IndexSearcher(reader);
        // 设置search,并且将search的策略设置为BM25
        searcher.setSimilarity(new BM25Similarity(k1, b));
        FileWriter fw_origin;
        FileWriter fw_expansion;
        if (null == gridResPath){
            fw_origin = new FileWriter(new File(this.runName + "-" + k1 + "-" + b + ".txt"));
            fw_expansion = new FileWriter(new File(this.runName + "-" + k1 + "-" + b + ".txt.expansion"));
            this.setResPath(this.runName + "-" + k1 + "-" + b + ".txt");
        }else{
            fw_origin = new FileWriter(new File(gridResPath));
            fw_expansion = new FileWriter(new File(gridResPath+".expansion"));
            this.setResPath(gridResPath);
        }

        int index = 0;
        for (String topic: topics){

            // original query 采用原始的query来检索
            System.out.println("origin query: " + topic);
            index += 1;
            String qid = new Integer(index).toString();
            String queryStr = DataPreprocess.remove_stemming(topic);
            TopDocs results = searcher.search(queryParser.parse(queryStr), this.resNum);
            ScoreDoc[] hists = results.scoreDocs;

            Document[] documents = new Document[100];
            for(int i = 0; i < relevantNum ; i ++){
                documents[i] = searcher.doc(hists[i].doc);
            }
            for(int i = 0; i < irrelevantNum; i++){
                documents[relevantNum+i] = searcher.doc(hists[hists.length - 1 - i].doc);
            }

            // 提取top1000写入origin中
            for (int i = 0; i < 1000; i++){
                Document doc = searcher.doc(hists[i].doc);
                //queryid Q0 docid rank score rankName
                String docPath = doc.get("path");
                String docID = docPath.substring(Math.max(docPath.lastIndexOf("/"),
                        docPath.lastIndexOf("\\")) + 1, docPath.lastIndexOf("."));
                String resStr = qid + " Q0 " + docID + " " + (i + 1) + " " + hists[i].score + " "+ this.runName + "\n";
                fw_origin.write(resStr);
            }

            // extent query 使用扩展后的query进行检索
            List<String> new_query_items = PRFTermDiscriminator.get_top_k_terms(relevantNum, irrelevantNum, documents, topic);
            StringBuffer buffer = new StringBuffer();
            for(String s: new_query_items){
                buffer.append(s);
                buffer.append(" ");
            }
            // 获得新的query的检索结果
            String extent_query = buffer.toString();
            System.out.println("extent query: " + extent_query);
            extent_query = DataPreprocess.remove_stemming(extent_query);
            results = searcher.search(queryParser.parse(extent_query), this.resNum);
            hists = results.scoreDocs;

            // 提取top1000写入expansion中
            for (int i = 0; i < 1000; i++){
                Document doc = searcher.doc(hists[i].doc);
                //queryid Q0 docid rank score rankName
                String docPath = doc.get("path");
                String docID = docPath.substring(Math.max(docPath.lastIndexOf("/"),
                        docPath.lastIndexOf("\\")) + 1, docPath.lastIndexOf("."));
                String resStr = qid + " Q0 " + docID + " " + (i + 1) + " " + hists[i].score + " "+ this.runName + "\n";
                fw_expansion.write(resStr);
            }
        }

        fw_origin.close();
        fw_expansion.close();
    }

    /**
     *
     * @param qrelsPath
     * @param gridResPath
     * @return
     */
    private float infNDCGOfQueryResult(String qrelsPath, String gridResPath) throws IOException {
        // used to exec in the command line
        String commandStr = "perl sample_eval.pl " + qrelsPath + " " + gridResPath;
        String res = Command.exeCmd(commandStr);
        return Float.parseFloat(res.split("\n")[1].split("\t")[4]);
    }

    public Entry<Float, Float> gridSearchBM25(String topicsPath,
                                              String qrelsPath, double[] kls, double[] bs) throws IOException, ParseException {
        String gridResPath = this.runName+"-grid-tmp.txt";
        double bestK1 = 1.2, bestB = 0.75, bestNDCG = -1;
        for (int i = 0; i < kls.length; i ++){
            for (int j = 0; j < bs.length; j++){
                // gridSearch逻辑部分
                this.QueryTopicBM25(topicsPath, (float)kls[i], (float)bs[j], gridResPath);
                float infNDCG = this.infNDCGOfQueryResult(qrelsPath, gridResPath);
                float infNDCG_expansion = this.infNDCGOfQueryResult(qrelsPath, gridResPath+".expansion");
                if (bestNDCG == 1 || bestNDCG < infNDCG_expansion){
                    bestNDCG = infNDCG_expansion;
                    bestK1 = kls[i];
                    bestB = bs[j];
                }
                System.out.println("current: K1="+kls[i] +" , b="+bs[j] + " ,infNDCG="+ infNDCG+
                        ",infNDCG_expansion="+ infNDCG_expansion + " ,bestNDCG=" + bestNDCG);
            }
        }
        System.out.println("===============================REPORT:===============================");
        System.out.println("best k1="+bestK1+", best b="+bestB+", bestNDCG="+bestNDCG);
        System.out.println("==========================Grid Search Over!==========================");
        return new AbstractMap.SimpleEntry((float)bestK1, (float)bestB);
    }

    public IndexReader getReader() {
        return reader;
    }

    public void setReader(IndexReader reader) {
        this.reader = reader;
    }

    public int getResNum() {
        return resNum;
    }

    public void setResNum(int resNum) {
        this.resNum = resNum;
    }

    public String getRunName() {
        return runName;
    }

    public void setRunName(String runName) {
        this.runName = runName;
    }

    public String getResPath() {
        return resPath;
    }

    public void setResPath(String resPath) {
        this.resPath = resPath;
    }
}
